import torch
import torch.nn as nn

"""
2D卷积；处理的是单帧/一张图像，通道独立。
处理一张(3, 224, 224)的图像，输入(1, 3, 224, 224) = (batch, channels, height, width)
对应的卷积核 (out_channel, in_channel, kh, kw)
"""
img_2d = torch.rand((1, 3, 224, 224))
conv2d = nn.Conv2d(
    in_channels=3,
    out_channels=64,
    kernel_size=(3, 3)
)
conv2d_img = conv2d(img_2d)
print(f"2D卷积后的图像形状：{conv2d_img.shape}")  # (1, 64, 222, 222)
print(f"2D卷积的权重：{conv2d.weight.shape}")  # (64, 3, 3, 3)
print(f"2D卷积的偏置：{conv2d.bias.shape}")  # (64,)

"""
3D卷积：处理的是连续帧或者体积数据，额外的维度表示深度或时间。
若处理一段16帧的CT扫描图像，输入形状(1, 1, 16, 256, 256) = (batch, channels, depth, height, width)
对应的卷积核 (out_channel, in_channel, kd, kh, kw)
"""
img_3d = torch.rand((1, 1, 16, 256, 256))
conv3d = nn.Conv3d(
    in_channels=1,
    out_channels=64,
    kernel_size=(3, 3, 3)
)
conv3d_img = conv3d(img_3d)
print(f"3D卷积后的图像形状：{conv3d_img.shape}")  # (1, 64, 14, 254, 254)
print(f"3D卷积的权重：{conv3d.weight.shape}")  # (64, 1, 3, 3, 3)
print(f"3D卷积的偏置：{conv3d.bias.shape}")  # (64)

'''
在不考虑偏置的情况
第一个2D卷积的参数量 = 64 * 3 * 3 * 3
第二个3D卷积的参数量 = 64 * 1 * 3 * 3 * 3
'''
